import json
import os
import os.path as osp
import argparse
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
import multiprocessing
from tictoc import TicToc


def read_camera_info(json_path):
    with open(json_path, 'r') as f:
        calibs = json.load(f)
    return calibs


def trans_img_new_intrin(out_intrin, new_size, old_intrin, img, filename):
    H = np.dot(out_intrin, np.linalg.inv(old_intrin))
    img = Image.fromarray(img[:, :, ::-1])

    # img.show()

    w, h = img.size
    H_inv = np.linalg.inv(H)
    H_inv = H_inv / H_inv[2, 2]  # 归一化
    # PIL expects a 3x3 matrix for perspective transform
    coefficients = H_inv.flatten()[:8]  # 取前8个元素，用于透视变换
    # 进行仿射变换
    if isinstance(new_size, list):
        w = new_size[1]
        h = new_size[0]
    new_img = img.transform(
        (w, h),
        Image.PERSPECTIVE,
        coefficients,
        Image.BICUBIC
    )
    # new_img = self.remove_black_borders(new_img)
    # self.show_diff_intrin_imgs(img, new_img)
    # img.show()
    # new_img.show()

    new_img.save(filename)
    return new_img


def main_worker(dataset_path, args):
    # cv2.getOptimalNewCameraMatrix(xxx,,,1)
    out_intrin = [
        [537.0237639534591,0.0,640.0],
        [0.0,537.0237639534591,360.0],
        [0.0,0.0,1.0]]
    # cv2.getOptimalNewCameraMatrix(xxx,,,0)
    out_intrin = [
        [570.0,0.0,640.0],
        [0.0,570.0,360.0],
        [0.0,0.0,1.0]]
    new_size = [720, 1280]
    json_path = osp.join(dataset_path, 'calib', 'calib.json')
    imgs_path = osp.join(dataset_path, 'camera')
    sensors_info = read_camera_info(json_path)
    select_cameras = ['camera75', 'camera77', 'camera80', 'camera81']
    for sensor_name in tqdm(select_cameras):
        assert sensor_name in sensors_info, f'{sensor_name} not in sensors_info'

        cam_data = sensors_info[sensor_name]
        # out_intrin = cam_data['normal']['K']
        # new_size = [cam_data['normal']['imgh'], cam_data['normal']['imgw']]

        img_path = osp.join(imgs_path, sensor_name)
        img_normal_path = osp.join(imgs_path, sensor_name + '_normal')
        if not os.path.exists(img_normal_path):
            os.mkdir(img_normal_path)
        img_files = os.listdir(img_path)
        for img_file in img_files:
            filename = osp.join(img_path, img_file)
            img = cv2.imread(filename)

            # cv2.imshow('ss',img)
            # cv2.waitKey()

            # disort >>>
            # 将RGB数组转换为BGR数组，因为OpenCV使用BGR格式
            # img = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
            # 假设你已经有了相机的内参矩阵cameraMatrix和畸变系数distCoeffs
            cameraMatrix = np.array(cam_data['ori']['K'])
            distCoeffs = np.array(cam_data['ori']['D'])

            # 获取图像尺寸
            h, w = img.shape[:2]
            img = cv2.resize(img, (cam_data['ori']['imgw'], cam_data['ori']['imgh']))
            # 获取新的相机矩阵（可选，如果需要裁剪图像）
            newcameramtx, roi = cv2.getOptimalNewCameraMatrix(cameraMatrix,
                                                              distCoeffs,
                                                              (cam_data['ori']['imgw'], cam_data['ori']['imgh']),
                                                              0)
            # 畸变校正
            # img = cv2.resize(img, (cam_data['imgw'], cam_data['imgh']))
            dst = cv2.undistort(img, cameraMatrix, distCoeffs, None, newcameramtx)
            # cv2.imshow('ss',dst)
            # cv2.waitKey()
            # 这里我们假设roi给出了一个包含整个有效视野的矩形，并且我们想要保持宽高比
            # 计算裁剪区域的坐标，以保持图像中心并尽可能接近1280x720（可能需要裁剪顶部或底部）
            # x = roi[0]
            # y = max(roi[1] - (cam_data['ori']['imgh'] - roi[3]) // 2, 0)  # 确保y坐标不会变成负数
            # width = roi[2]
            # height = min(roi[3], cam_data['ori']['imgh'])  # 确保高度不会超过目标高度
            # # 裁剪图像
            # cropped_dst = dst[y:y + height, x:x + width]

            # cv2.imshow('ss',cropped_dst)
            # cv2.waitKey()

            # # 如果裁剪后的图像不是1280x720，我们可能需要进一步缩放它
            # if cropped_dst.shape[1] != cam_data['ori']['imgw'] or cropped_dst.shape[0] != cam_data['ori']['imgh']:
            #     cropped_dst = cv2.resize(cropped_dst, (cam_data['ori']['imgw'], cam_data['ori']['imgh']))
            # dst = cropped_dst

            # cv2.namedWindow('ss', cv2.WINDOW_NORMAL)
            # cv2.imshow('ss',dst)
            # cv2.waitKey()
            # dst = cv2.resize(dst, (w, h))

            # cv2.imshow('ori',img)
            # cv2.imshow('dst',dst)
            # cv2.waitKey()
            trans_img_new_intrin(out_intrin, new_size, newcameramtx, dst, osp.join(img_normal_path, img_file))

            # output_file = osp.join(img_normal_path,img_file)
            # cv2.imwrite(output_file, dst)

        with open(json_path, 'r') as f:
            data = json.load(f)

        data[sensor_name]['normal'] = dict()
        data[sensor_name]['normal']['imgh'] = new_size[0]
        data[sensor_name]['normal']['imgw'] = new_size[1]
        data[sensor_name]['normal']['K'] = [out_intrin[i] for i in range(3)]

        with open(json_path, 'w') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f'{sensor_name} done')


def main(args):
    # 统计耗时
    cost = TicToc("相机归一化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/media/adt/T7/ZWH/docker/files/data/cyw2",
                        help='your data root for kitti')
    args = parser.parse_args()

    main(args)
